from .exceptions import IndexErrorOfStructure


class SegmentTree:
    """使用列表按照完全二叉树的方法存储线段树，如果元素数为n，最多需要4O(n)的空间"""
    def __init__(self, arr, merger):
        # merger定义期间的行为，或者说如何通过左右子树的区间的行为构建上层的区间
        self._data = list(arr)
        self.merger = merger
        self._tree = [None for _ in range(4 * len(arr))]
        self._build_tree(0, 0, len(self._data) - 1)

    def _build_tree(self, index, begin, end):
        """
        创建线段树，在index位置创建表示[begin, end]的线段树
        @index 节点索引
        @begin 区间左端点
        @end 区间右端点
        """
        if begin == end:
            self._tree[index] = self._data[begin]
            return
        left_index = self._left(index)
        right_index = self._right(index)
        mid = (begin + end) // 2
        self._build_tree(left_index, begin, mid)
        self._build_tree(right_index, mid+1, end)
        self._tree[index] = self.merger(self._tree[left_index], self._tree[right_index])

    def size(self):
        return len(self._data)

    def get(self, index):
        if index < 0 or index >= len(self._data):
            raise IndexErrorOfStructure("Invalid index.")
        return self._data[index]

    def _left(self, index):
        """返回节点左孩子节点的索引"""
        return 2 * index + 1

    def _right(self, index):
        """返回节点右孩子节点的索引"""
        return 2 * index + 2

    def query(self, query_l, query_r):
        length = len(self._data)
        if query_l < 0 or query_l >= length or query_r < 0 or query_r >= length or query_l > query_r:
            raise IndexErrorOfStructure("Invalid index.")
        return self._query(0, 0, length - 1, query_l, query_r)

    def _query(self, tree_index, l, r, query_l, query_r):
        if l == query_l and r == query_r:
            return self._tree[tree_index]
        mid = (l + r) // 2
        l_child_index = self._left(tree_index)
        r_right_index = self._right(tree_index)
        if query_l >= mid + 1:
            return self._query(r_right_index, mid+1, r, query_l, query_r)
        if query_r <= mid:
            return self._query(l_child_index, l, mid, query_l, query_r)
        left_result = self._query(l_child_index, l, mid, query_l, mid)
        right_result = self._query(r_right_index, mid+1, r, mid+1, query_r)
        return self.merger(left_result, right_result)

    def set(self, index, e):
        if index < 0 or index >= len(self._data):
            raise IndexErrorOfStructure("Invalid index.")
        self._data[index] = e
        self._set(0, 0, len(self._data) - 1, index, e)

    def _set(self, tree_index, l, r, index, e):
        if l == r:
            self._tree[tree_index] = e
            return
        mid = (l + r) // 2
        l_child_index = self._left(tree_index)
        r_child_index = self._right(tree_index)
        if index >= mid + 1:
            self._set(r_child_index, mid+1, r, index, e)
        else:
            self._set(l_child_index, l, mid, index, e)
        # 记得更新每一层的值
        self._tree[tree_index] = self.merger(self._tree[l_child_index], self._tree[r_child_index])